# @Time : 2021/8/4 10:05
# @Author : Li Kunlun
# @Description : softmax回归的从零开始实现
import utils as d2l
from mxnet import autograd, nd

# 1、读取数据集
batch_size = 256  # 使用fashion-mnist数据集，设置批量大小为256
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size)

# 2、初始化模型
num_inputs = 784  # 模型输入向量长度28*28=784
num_outputs = 10  # 输出10个类别

W = nd.random.normal(scale=0.01, shape=(num_inputs, num_outputs))
b = nd.zeros(num_outputs)

# 为模型参数附上梯度
W.attach_grad()
b.attach_grad()

# 3、实现softmax运算
"""
1、对多维NDArray操作
    [[ 5.  7.  9.]]
    <NDArray 1x3 @cpu(0)> 
    [[  6.]
     [ 15.]]
    <NDArray 2x1 @cpu(0)>
"""


# X = nd.array([[1, 2, 3], [4, 5, 6]])
# print(X.sum(axis=0, keepdims=True), X.sum(axis=1, keepdims=True))

def softmax(X):
    X_exp = X.exp()
    partition = X_exp.sum(axis=1, keepdims=True)
    return X_exp / partition  # 这里应用了广播机制


"""
[[ 0.62647122  0.126293    0.01826552  0.10885343  0.12011679]
 [ 0.25569436  0.2917251   0.07546549  0.30240679  0.07470828]]
<NDArray 2x5 @cpu(0)> 
[ 0.99999994  1.        ]
<NDArray 2 @cpu(0)>
"""


# X = nd.random.normal(shape=(2, 5))
# X_prob = softmax(X)
# print(X_prob, X_prob.sum(axis=1))


# 4、定义模型
def net(X):
    """
    :param
        :param X: 28*28大小的图片
        :return: (1*10)输出
    """
    return softmax(nd.dot(X.reshape((-1, num_inputs)), W) + b)  # X转换成num_inputs列（1*784）


# 5、定义损失函数

# 交叉熵函数
def cross_entropy(y_hat, y):
    return -nd.pick(y_hat, y).log()


# 6、计算分类准确率
def accuracy(y_hat, y):
    return (y_hat.argmax(axis=1) == y.astype('float32')).mean().asscalar()


# 评价模型net在数据集data_iter上的准确率
def evaluate_accuracy(data_iter, net):
    acc_sum, n = 0.0, 0
    for X, y in data_iter:
        y = y.astype('float32')
        acc_sum += (net(X).argmax(axis=1) == y).sum().asscalar()
        n += y.size
    return acc_sum / n


# evaluate_accuracy(test_iter, net)

# 7、训练模型
num_epochs, lr = 5, 0.1


def train_ch3(net, train_iter, test_iter, loss, num_epochs, batch_size, params=None, lr=None, trainer=None):
    """
    :param
        :param net: 定义模型 (1*10)输出
        :param train_iter:
        :param test_iter:
        :param loss: 交叉熵
        :param num_epochs:训练周期5
        :param batch_size: 批量处理256
        :param params: [W, b]
        :param lr: 学习率
        :param trainer:
        :return:
    """
    for epoch in range(num_epochs):
        train_l_sum, train_acc_sum, n = 0.0, 0.0, 0
        for X, y in train_iter:
            with autograd.record():
                y_hat = net(X)
                l = loss(y_hat, y).sum()
            l.backward()
            if trainer is None:
                d2l.sgd(params, lr, batch_size)
            else:
                trainer.step(batch_size)  # “softmax回归的简洁实现”一节将用到
            y = y.astype('float32')
            train_l_sum += l.asscalar()
            train_acc_sum += (y_hat.argmax(axis=1) == y).sum().asscalar()
            n += y.size
        test_acc = evaluate_accuracy(test_iter, net)
        print('epoch %d, loss %.4f, train acc %.3f, test acc %.3f'
              % (epoch + 1, train_l_sum / n, train_acc_sum / n, test_acc))


train_ch3(net, train_iter, test_iter, cross_entropy, num_epochs, batch_size, [W, b], lr)

# 8、预测
for X, y in test_iter:
    break

true_labels = d2l.get_fashion_mnist_labels(y.asnumpy())
pred_labels = d2l.get_fashion_mnist_labels(net(X).argmax(axis=1).asnumpy())
titles = [true + '\n' + pred for true, pred in zip(true_labels, pred_labels)]

d2l.show_fashion_mnist(X[0:9], titles[0:9])
